import gradio as gr
import os
import glob
from datetime import datetime
import base64
import io
import tempfile
import speech_recognition as sr
from gtts import gTTS
import time
import uuid
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

from langchain_nvidia_ai_endpoints import ChatNVIDIA, NVIDIAEmbeddings
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain.schema import HumanMessage, AIMessage
from langchain.prompts import MessagesPlaceholder
from langchain_community.vectorstores import FAISS
from langchain.text_splitter import CharacterTextSplitter

# 设置 API 密钥
NVIDIA_API_KEY = os.getenv('NVIDIA_API_KEY')
# 初始化嵌入模型
embedder = NVIDIAEmbeddings(model="NV-Embed-QA")

# 加载 robot 嵌入
def load_robot_embeddings(embedding_dir="./zh_data/embeddings"):
    """
    加载最新的机器人嵌入
    
    参数:
        embedding_dir: 嵌入文件目录
        embedder: 嵌入模型，如果为None则使用默认模型
    
    返回:
        向量存储对象
    """
    # 确保目录存在
    os.makedirs(embedding_dir, exist_ok=True)
    
    # 获取最新的嵌入文件
    embedding_files = glob.glob(os.path.join(embedding_dir, "robot_embedding_*"))
    if not embedding_files:
        print("未找到机器人嵌入文件")
        return None
    
    # 按时间戳排序，获取最新的嵌入文件
    latest_embedding = max(embedding_files, key=os.path.getctime)
    print(f"加载最新的机器人嵌入: {latest_embedding}")
    
    # 加载向量存储
    try:
        store = FAISS.load_local(latest_embedding, embedder, allow_dangerous_deserialization=True)
        print("机器人嵌入加载成功")
        return store
    except Exception as e:
        print(f"加载机器人嵌入时出错: {str(e)}")
        return None

# 加载商品嵌入
def load_product_embeddings(embedding_dir="./zh_data/embeddings"):
    """
    加载最新的商品嵌入
    
    参数:
        embedding_dir: 嵌入文件目录
        embedder: 嵌入模型，如果为None则使用默认模型
    
    返回:
        向量存储对象
    """
    # 确保目录存在
    os.makedirs(embedding_dir, exist_ok=True)
    
    # 获取最新的嵌入文件
    embedding_files = glob.glob(os.path.join(embedding_dir, "product_embedding_*"))
    if not embedding_files:
        print("未找到商品嵌入文件")
        return None
    
    # 按时间戳排序，获取最新的嵌入文件
    latest_embedding = max(embedding_files, key=os.path.getctime)
    print(f"加载最新的商品嵌入: {latest_embedding}")
    
    # 加载向量存储
    try:
        store = FAISS.load_local(latest_embedding, embedder, allow_dangerous_deserialization=True)
        print("商品嵌入加载成功")
        return store
    except Exception as e:
        print(f"加载商品嵌入时出错: {str(e)}")
        return None

# 加载嵌入
robot_store = load_robot_embeddings()
product_store = load_product_embeddings()

if robot_store:
    print("成功加载 robot 嵌入，可用于对话问答")
else:
    print("未能加载 robot 嵌入，将使用默认嵌入")

if product_store:
    print("成功加载商品嵌入，可用于商品检索")
else:
    print("未能加载商品嵌入，将使用默认嵌入")

# 打印可用模型列表
# print("正在获取可用模型列表...")
# try:
#     available_models = ChatNVIDIA.get_available_models()
# except Exception as e:
#     print(f"获取模型列表时出错: {str(e)}")
#     print("请检查您的 API 密钥是否正确")

# 初始化聊天模型
try:
    llm = ChatNVIDIA(
        model="deepseek-ai/deepseek-r1",
        # model="meta/llama-3.1-405b-instruct",
        # model="microsoft/phi-3.5-vision-instruct",
        # model="nvidia/neva-22b",
        model_type="chat",
    )
    print("模型初始化成功")
except Exception as e:
    print(f"模型初始化失败: {str(e)}")
    print("请检查模型名称和 API 密钥是否正确")
    raise

# 创建提示模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个专业富有亲和力的导购，销售服务机器人。客户有任何问题，你都可以回答。客户提出的问题比较笼统时，你要追问挖掘细节关键需求，推荐合适的产品。请记住之前的对话内容，保持上下文的连贯性。"),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{question}")
])

# 创建带上下文的提示模板
context_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个专业富有亲和力的导购，销售商品。客户有任何问题，你都可以回答。客户提出的问题比较笼统时，你要追问挖掘细节关键需求，推荐合适的产品。请记住之前的对话内容，保持上下文的连贯性。\n\n以下是相关的上下文信息，请根据这些信息回答问题：\n{context}"),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{question}")
])

# 创建 robot 提示模板
robot_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个专业的机器人助手，擅长回答关于机器人的各种问题。请根据提供的上下文信息回答问题，如果上下文中没有相关信息，请诚实地告诉用户你不知道。请记住之前的对话内容，保持上下文的连贯性。\n\n以下是相关的上下文信息，请根据这些信息回答问题：\n{context}"),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{question}")
])

# 创建图片分析提示模板
image_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个专业的商品导购助手，擅长分析图片中的商品并推荐相似产品。请根据提供的图片和上下文信息，详细描述图片中的商品，并推荐相似的产品。如果上下文中没有相关信息，请诚实地告诉用户你不知道。请记住之前的对话内容，保持上下文的连贯性。\n\n以下是相关的上下文信息，请根据这些信息回答问题：\n{context}"),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{question}")
])

def format_chat_history(history):
    """格式化聊天历史"""
    formatted_history = []
    for human_msg, assistant_msg in history:
        # 跳过空消息
        if human_msg == "" and assistant_msg.startswith("👋 您好！我是智购助手"):
            continue
        formatted_history.extend([
            HumanMessage(content=human_msg),
            AIMessage(content=assistant_msg)
        ])
    return formatted_history

def chat_with_bot(message, history, use_robot=False, use_local_products=False, image=None):
    """处理聊天消息的函数"""
    global llm  # 将global声明移到函数开头
    
    try:
        # 检查llm是否已初始化
        if 'llm' not in globals() or llm is None:
            # 如果llm未初始化，尝试初始化
            try:
                llm = ChatNVIDIA(
                    # model="deepseek-ai/deepseek-r1",
                    model="meta/llama-3.1-405b-instruct",
                    model_type="chat",
                )
                print("模型初始化成功")
            except Exception as e:
                error_message = f"模型初始化失败: {str(e)}\n请检查模型名称和 API 密钥是否正确。"
                history.append((message, error_message))
                return history
        
        # 如果是欢迎语，直接返回
        if message is None and len(history) == 1 and "👋 您好！我是智购助手" in history[0][1]:
            return history
        
        # 格式化聊天历史
        chat_history = format_chat_history(history)
        
        # 选择使用的向量存储和提示模板
        if image is not None:
            # 如果有图片，使用图片分析模式
            current_store = product_store if product_store else None
            current_prompt = image_prompt
            print("使用图片分析模式")
            
            # 将图片转换为 base64 字符串
            if isinstance(image, str) and image.startswith('data:image'):
                # 已经是 base64 格式
                image_base64 = image
            else:
                # 将图片转换为 base64
                buffered = io.BytesIO()
                
                # 调整图片大小，限制最大尺寸为 800x800
                if hasattr(image, 'size'):
                    width, height = image.size
                    
                    # 计算当前图片的 token 数量（粗略估计）
                    # 每个像素大约需要 1.5 个 token
                    current_tokens = width * height * 1.5
                    max_tokens = 1024
                    
                    if current_tokens > max_tokens:
                        print(f"图片 token 数量 ({current_tokens:.0f}) 超过限制 ({max_tokens})，正在压缩...")
                        # 计算压缩比例
                        ratio = (max_tokens / current_tokens) ** 0.5
                        new_width = int(width * ratio)
                        new_height = int(height * ratio)
                        image = image.resize((new_width, new_height), Image.LANCZOS)
                        print(f"图片已压缩至: {new_width}x{new_height}")
                
                image.save(buffered, format="JPEG", quality=70)  # 使用JPEG格式和较低质量以减小文件大小
                img_str = base64.b64encode(buffered.getvalue()).decode()
                image_base64 = f"data:image/jpeg;base64,{img_str}"
            
            # 在消息中添加图片信息
            message = f"{message}\n[图片内容: {image_base64}]"
        elif use_robot and robot_store:
            current_store = robot_store
            current_prompt = robot_prompt
            print("使用 robot 嵌入进行回答")
        elif use_local_products:
            # 检查product_store是否存在
            if product_store is None:
                print("警告: 本地商品知识库未加载，将使用默认嵌入进行回答")
                current_store = None
                current_prompt = prompt
            else:
                current_store = product_store
                current_prompt = context_prompt
                print("使用本地商品知识库进行回答")
        else:
            # 使用默认的向量存储，如果没有定义则使用product_store
            current_store = None
            current_prompt = prompt
            print("使用默认嵌入进行回答")
        
        # 创建处理链
        if current_store is not None:
            # 如果有向量存储，使用检索增强生成
            try:
                retriever = current_store.as_retriever()
                if retriever is not None:
                    print(f"检索器有效")
                    chain = (
                        {
                            "context": retriever,
                            "question": lambda x: x,
                            "history": lambda x: chat_history
                        }
                        | current_prompt
                        | llm
                        | StrOutputParser()
                    )
                    print(f"chain finished")
                else:
                    print(f"检索器为空")
            except Exception as e:
                print(f"创建检索器时出错: {str(e)}")
                # 如果创建检索器失败，回退到默认模式
                chain = (
                    {
                        "question": lambda x: x,
                        "history": lambda x: chat_history
                    }
                    | prompt
                    | llm
                    | StrOutputParser()
                )
        else:
            # 如果没有向量存储，直接使用模型回答
            chain = (
                {
                    "question": lambda x: x,
                    "history": lambda x: chat_history
                }
                | current_prompt
                | llm
                | StrOutputParser()
            )
        
        # 调用模型获取响应
        response = chain.invoke(message)
        
        # 更新聊天历史
        history.append((message, response))
        
        # 返回更新后的历史
        return history
    except Exception as e:
        error_message = f"发生错误: {str(e)}\n\n可能的原因:\n1. API 密钥无效或过期\n2. 模型访问权限不足\n3. 网络连接问题\n\n请检查您的 API 密钥和网络连接。"
        history.append((message, error_message))
        return history

def speech_to_text(audio):
    """将语音转换为文本"""
    if audio is None:
        return None
    
    try:
        # 初始化识别器
        recognizer = sr.Recognizer()
        
        # 从音频文件中读取音频
        with sr.AudioFile(audio) as source:
            audio_data = recognizer.record(source)
        
        # 使用 Google 语音识别
        text = recognizer.recognize_google(audio_data, language='zh-CN')
        print(f"语音识别结果: {text}")
        return text
    except Exception as e:
        print(f"语音识别出错: {str(e)}")
        return None

def text_to_speech(text):
    """将文本转换为语音"""
    if text is None or text == "":
        return None
    
    try:
        # 如果输入是列表，取最后一条消息
        if isinstance(text, list) and len(text) > 0:
            # 获取最后一条消息的回复部分
            last_message = text[-1][1]
            text = last_message
        
        # 创建临时文件
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp3')
        temp_filename = temp_file.name
        temp_file.close()
        
        # 使用 gTTS 将文本转换为语音
        tts = gTTS(text=text, lang='zh-CN')
        tts.save(temp_filename)
        
        return temp_filename
    except Exception as e:
        print(f"文本转语音出错: {str(e)}")
        return None

# 创建 Gradio 界面
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown("""
    # 🛍️ 智购助手
    
    *智能导购，让购物更简单*
    """)
    
    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(
                height=500,
                value=[("", "👋 您好！我是智购助手，很高兴为您服务！\n\n我可以帮您：\n1️⃣ 推荐合适的商品\n2️⃣ 分析商品图片\n3️⃣ 回答购物相关问题\n\n请问有什么可以帮您的吗？")],
                bubble_full_width=False,
                show_copy_button=True,
                avatar_images=(
                    "web_source/user.jpg",  # 用户头像
                    "web_source/robot.jpg"   # 机器人头像
                ),
                show_label=False
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    label="",
                    placeholder="请输入您的问题，例如：我想买一件衣服...",
                    lines=2,
                    container=False,
                    scale=8
                )
                submit = gr.Button("发送", variant="primary", scale=1)
                clear = gr.Button("清除", variant="secondary", scale=1)
            
            with gr.Accordion("高级功能", open=False):
                with gr.Row():
                    audio_input = gr.Audio(
                        sources=["microphone"], 
                        type="filepath", 
                        label="语音输入",
                        show_label=True
                    )
                    audio_output = gr.Audio(
                        label="语音输出",
                        type="filepath",
                        show_label=True
                    )
                
                with gr.Row():
                    image_input = gr.Image(
                        label="上传商品图片",
                        type="pil",
                        show_label=True,
                        height=200
                    )
                    image_submit = gr.Button("分析图片", variant="primary")
            
            with gr.Row():
                with gr.Column(scale=1):
                    use_robot = gr.Checkbox(
                        label="使用机器人知识库",
                        value=False,
                        info="使用预定义的机器人知识库"
                    )
                with gr.Column(scale=1):
                    use_local_products = gr.Checkbox(
                        label="使用本地商品库",
                        value=False,
                        info="使用本地商品知识库"
                    )
    
    # 设置事件处理
    msg.submit(
        chat_with_bot,
        [msg, chatbot, use_robot, use_local_products, image_input],
        [chatbot],
        queue=True
    ).then(
        lambda: "", None, [msg]
    )
    
    submit.click(
        chat_with_bot,
        [msg, chatbot, use_robot, use_local_products, image_input],
        [chatbot],
        queue=True
    ).then(
        lambda: "", None, [msg]
    )
    
    clear.click(lambda: [], None, chatbot)
    
    image_submit.click(
        chat_with_bot,
        [msg, chatbot, use_robot, use_local_products, image_input],
        [chatbot],
        queue=True
    )
    
    # 添加音频处理
    # audio_input.stop_recording(
    #     transcribe_audio,
    #     [audio_input, chatbot, use_robot, use_local_products, image_input],
    #     [chatbot, audio_output]
    # ).then(
    #     lambda: "", None, [msg]
    # )

# 启动界面
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        favicon_path="./favicon.ico" if os.path.exists("./favicon.ico") else None
    )
